feat: opt-in pipeline fit optimisations, native bucketize, and sampled fitting support with caching by ConorWorthington · Pull Request #69 · ExpediaGroup/kamae

ConorWorthington · 2026-08-05T13:56:38Z

Description

Provide a short description of the PR changes.

The below checklists come from the docs page on adding new transformers here

Keras Layer Checklist

Verify that:

The new Keras layer extends BaseLayer
The _call method has been implemented in the new layer.
The compatible_dtypes property is defined in the new layer.
The new layer is decorated with @tf.keras.utils.register_keras_serializable(package=kamae.__name__).
The new layer takes a name, input_dtype, and output_dtype as arguments to the constructor and that this is passed to the super constructor.
The Keras layer is serializable. I have implemented the get_config method.
There are unit tests of the new layer.
There is a specific test of layer serialisation added here.
The new layer is imported in the init.py file in the layers directory.

Spark Transformer/Estimator Checklist

Verify that:

The new Spark Transformer extends BaseTransformer.
If the new transform needs a fit method, a Spark Estimator has been implemented that extends BaseEstimator.
The instructions in the above docs page have been followed for the __init__ and setParams methods.
The transformer uses one of the input/output mixin classes from base.py.
If the new transformer requires more parameters that would need to be serialised to the Spark ML pipeline, there is a implemented parameter class by extending the Params class here.
The compatible_dtypes property has been implemented to specify the input/output data types that my transformer/estimator supports.
A Keras subclassed layer is returned in the transformer's get_tf_layer method.
There are unit tests of the new transform. In particular, there are parity tests between the Spark and Keras implementations.
The new transformer/estimator is imported in the init.py file in the transformers/estimators directory.

Finally, please verify that:

There is a new entry (alphabetical order) in the README table describing the new layer/transformer

Plans are too slow to materialise - this is our attempt to speed it up

Wrap the moments aggregation in StandardScale, SingleFeatureArrayStandardScale and ConditionalStandardScale estimators in a guarded persist/unpersist so the array-size probe and the aggregation reuse a materialised result instead of re-scanning the upstream lineage twice. Repair the incomplete persist edit in ConditionalStandardScale._fit. Add checkpointInterval / pruneInputColumns coverage to the pipeline tests and a checkpoint directory to the spark_session fixture. Surface estimator fit errors as RuntimeError chained from the original exception. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Merge latest version

georyetti

Mostly minor changes to pipeline logic. Also bucketize changes I assume should not be here. Lastly you need to run a uv lock.

georyetti · 2026-08-12T09:19:03Z

        """
-        super().__init__(stages=stages)
+        kwargs = self._input_kwargs
+        super().__init__()


The new super call drops the stages=stages, was this intended?

georyetti · 2026-08-12T09:19:31Z

+        kwargs = self._input_kwargs
+        super().__init__()
+        self._setDefault(
+            checkpointInterval=0,


Default to None instead of 0?

georyetti · 2026-08-12T09:25:26Z

+            for param in stage.params:
+                if not (param.name.endswith("Col") or param.name.endswith("Cols")):
+                    continue
+                if not stage.isDefined(param):
+                    continue
+                value = stage.getOrDefault(param)
+                if isinstance(value, str):
+                    required_input_columns.add(value)
+                elif isinstance(value, (list, tuple)):
+                    required_input_columns.update(
+                        item for item in value if isinstance(item, str)
+                    )


This feels more complex than I think it needs to be. Two things:

Why do we need to iterate through the stage.params at all if we have already added the inputs using get_layer_inputs_outputs

Even if we do need, we can just check for presence of inputCol or inputCols as a stage always has strictly one of these defined. if stage.hasParam("inputCol") and stage.isDefined("inputCol"): value = stage.getInputCol() and after check input cols. But this is really what get_layer_inputs_outputs does.

You're right that inputCol/inputCols are already covered by get_layer_inputs_outputs but the sweep isn't for those.

It's there for the other column params a stage reads that aren't its input col: maskCols/relevanceCol on ConditionalStandardScaleEstimator, and queryIdCol on the listwise transformers. Those get read at fit time, so if we don't keep them, pruneInputColumns=True drops them and the fit blows up with "column not found". There are regression tests covering exactly that.

If the suffix-matching feels too hacky, I'm happy to swap it for a small get_fit_input_columns() hook on the handful of stages that need it instead. Let me know what you'd prefer.

georyetti · 2026-08-12T09:26:56Z

+        """
+        required_input_columns = self.collect_required_input_columns(stages)
+        columns_to_keep = [c for c in dataset.columns if c in required_input_columns]
+        if columns_to_keep and len(columns_to_keep) < len(dataset.columns):


Pedantic but do we need that second condition? If we are inside this function then we are pruning, and columns_to_keep is always a subset. So I would just check it's not empty and otherwise select

georyetti · 2026-08-12T09:28:00Z

+        fit. 0 (or None) disables checkpointing.
+        :returns: KamaeSparkPipeline object with checkpointInterval set.
+        """
+        return self._set(checkpointInterval=value)


We should error on checkpoint interval being negative here. Personally I think we should error on 0 too and treat None as the no checkpoint behaviour

georyetti · 2026-08-12T09:29:50Z

+        :returns: KamaeSparkPipeline object with params set.
        """
        kwargs = self._input_kwargs
        return self._set(**kwargs)


This set params does not use the setter methods at all. So any validation in them will not be respected when the user passes arguments to the init as opposed to using the setter method. If you check how I defined the setParams for the estimator and transformer you can see I use the setter method.

georyetti · 2026-08-12T11:11:43Z

-                return None
-            # We add 1 because we want to reserve the 0 index for mask/padding.
-            return bisect_right(splits, value) + 1
+        def bucketize(value: Column) -> Column:


Same to assume that all this bucket logic changes to this transformer are not meant to be in this PR?

I was looking for inefficiencies in the library... I was going to patch more but only bucketize was really impacted in a way I could speed up simply. We don't really use it but figured it was a good improvement so may as well include.

georyetti · 2026-08-12T11:12:50Z

 dependencies = [
    "pyspark>=3.4.0,<4.0.0",
    "pandas>=1.3.4,<3.0.0",
+    "pyarrow>=4.0.0",


Adding new dependencies needs a uv lock pls

Lock updated

- Restore stages=stages in __init__ super call - checkpointInterval defaults to None; reject non-positive via setter - Route setParams through setter methods so validation runs - Drop redundant length check in prune_unused_input_columns - Regenerate uv.lock to include pyarrow (required by pandas_udf) Retains the aux-column sweep in collect_required_input_columns: it is load-bearing for pruning correctness (maskCols/relevanceCol/queryIdCol are not returned by get_layer_inputs_outputs) and defended by regression tests. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

ConorWorthington and others added 24 commits June 29, 2026 16:31

Added localCheckPointInterval

96bc515

Plans are too slow to materialise - this is our attempt to speed it up

Update pipeline.py

b581ff1

Update single_feature_array_standard_scale.py

32502cc

Update standard_scale.py

4e367d3

Update conditional_standard_scale.py

2d2f75a

Update base.py

f16c0ef

Update bucketize.py

8089c4f

Update bucketize.py

9453119

Update conditional_standard_scale.py

06043b5

Update single_feature_array_standard_scale.py

b1096ce

Update standard_scale.py

70cd178

Update bucketize.py

c3db6b7

Update base.py

847f181

Update pipeline.py

6d5964a

Update pipeline.py

ff3e9f7

Update standard_scale.py

f8cff4d

Update single_feature_array_standard_scale.py

5bb8cd9

Update conditional_standard_scale.py

4c470d4

Update transform_utils.py

a8a8750

Update pyproject.toml

71cf01a

Update pipeline.py

6c4e18f

Update pipeline.py

771ab9c

Merge pull request #1 from ConorWorthington/pipeline-fit-optimisations

a2985e1

Merge latest version

ConorWorthington requested a review from a team as a code owner August 5, 2026 13:56

ConorWorthington requested review from ddonghi and mandrecki August 5, 2026 13:56

ConorWorthington changed the title ~~feat: opt-in pipeline fit optimisations, native bucketize, and sampled fitting~~ feat: opt-in pipeline fit optimisations, native bucketize, and sampled fitting support with caching Aug 5, 2026

ConorWorthington added 2 commits August 5, 2026 16:56

Patch black

94dea39

Patch error

4f5cea2

ConorWorthington added 4 commits August 5, 2026 16:57

Patch test case

3878d2f

Remove complexity, remove comments

1918090

Patch for lint

96fc184

Patch for linter

6af4ba4

georyetti requested changes Aug 12, 2026

View reviewed changes

Navigation Menu

Uh oh!

feat: opt-in pipeline fit optimisations, native bucketize, and sampled fitting support with caching - #69

feat: opt-in pipeline fit optimisations, native bucketize, and sampled fitting support with caching#69
ConorWorthington wants to merge 31 commits into
ExpediaGroup:mainfrom
ConorWorthington:main

ConorWorthington commented Aug 5, 2026

Uh oh!

georyetti left a comment

Uh oh!

georyetti Aug 12, 2026

Uh oh!

georyetti Aug 12, 2026

Uh oh!

georyetti Aug 12, 2026

Uh oh!

ConorWorthington Aug 12, 2026

Uh oh!

georyetti Aug 12, 2026

Uh oh!

georyetti Aug 12, 2026

Uh oh!

ConorWorthington Aug 12, 2026

Uh oh!

georyetti Aug 12, 2026

Uh oh!

georyetti Aug 12, 2026

Uh oh!

ConorWorthington Aug 12, 2026

Uh oh!

georyetti Aug 12, 2026

Uh oh!

ConorWorthington Aug 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

ConorWorthington commented Aug 5, 2026

Description

Keras Layer Checklist

Spark Transformer/Estimator Checklist

Uh oh!

georyetti left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants